'''
三分类模型：CT和PET组成双通道输入，在全连接层与额外信息连接
'''
import numpy as np
import os
import cv2
import time
import shutil
import datetime
import pandas as pd
from glob import glob
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.metrics import classification_report, confusion_matrix
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.callbacks import TensorBoard, ReduceLROnPlateau, ModelCheckpoint
from tensorflow.keras.optimizers import SGD, Adam
from tensorflow.keras.regularizers import l2
from tensorflow.keras.initializers import glorot_uniform
from tensorflow.keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D
from tensorflow.keras.layers import AveragePooling2D, MaxPooling2D, Concatenate, Dropout

model_name = "three_resnet50_ct_pet_1"
time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
log_dir = "/home/zmy/workspace2021/three_logs/" + model_name + time_str
work_dir = "/home/zmy/workspace2021/three_workdirs/" + model_name + time_str
model_save_path = "/home/zmy/workspace2021/three_models/" + model_name
Batch_size = 8
EPOCH = 100
Learn_rate = 0.0001

def identity_block(X, f, filters, stage, block):
    """
    Implementation of the identity block as defined in Figure 3

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network

    Returns:
    X -- output of the identity block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value. You'll need this later to add back to the main path.
    X_shortcut = X

    # First component of main path
    X = Conv2D(filters=F1, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2a',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def convolutional_block(X, f, filters, stage, block, s=2):
    """
    Implementation of the convolutional block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network
    s -- Integer, specifying the stride to be used

    Returns:
    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value
    X_shortcut = X

    ##### MAIN PATH #####
    # First component of main path
    X = Conv2D(F1, (1, 1), strides=(s, s), name=conv_name_base + '2a', kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    ##### SHORTCUT PATH #### (≈2 lines)
    X_shortcut = Conv2D(filters=F3, kernel_size=(1, 1), strides=(s, s), padding='valid', name=conv_name_base + '1',
                        kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    X_shortcut = BatchNormalization(axis=3, name=bn_name_base + '1')(X_shortcut)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def ResNet50(input_shape=(512, 512, 2), other_info=(13, ), classes=3):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """

    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)
    other_info = Input(shape=other_info)
    # Zero-Padding
    X = ZeroPadding2D((3, 3))(X_input)

    # Stage 1
    X = Conv2D(64, (7, 7), strides=(2, 2), name='conv1', kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name='bn_conv1')(X)
    X = Activation('relu')(X)
    X = MaxPooling2D((3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, f=3, filters=[64, 64, 256], stage=2, block='a', s=1)
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='b')
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='c')

    ### START CODE HERE ###

    # Stage 3 (≈4 lines)
    X = convolutional_block(X, f=3, filters=[128, 128, 512], stage=3, block='a', s=2)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='b')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='c')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='d')

    # Stage 4 (≈6 lines)
    X = convolutional_block(X, f=3, filters=[256, 256, 1024], stage=4, block='a', s=2)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='b')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='c')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='d')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='e')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='f')

    # Stage 5 (≈3 lines)
    X = convolutional_block(X, f=3, filters=[512, 512, 2048], stage=5, block='a', s=2)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='b')
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='c')

    # AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    X = AveragePooling2D((2, 2), name="avg_pool")(X)

    ### END CODE HERE ###

    # output layer
    X = Flatten()(X)
    X = Concatenate()([X, other_info])
    X = Dense(512, kernel_regularizer=l2(0.001), activation='relu', kernel_initializer=glorot_uniform(seed=0))(X)
    X = Dropout(0.5)(X)
    X = Dense(512, kernel_regularizer=l2(0.001), activation='relu', kernel_initializer=glorot_uniform(seed=0))(X)
    X = Dense(classes, activation='softmax', name='fc' + str(classes), kernel_initializer=glorot_uniform(seed=0))(X)

    # Create model
    model = Model(inputs=[X_input, other_info], outputs=X, name=model_name)
    return model


def standard_scaler(features):
    scaler = MinMaxScaler()
    # scaler = StandardScaler()
    x_train = scaler.fit_transform(features)
    
    return x_train

def read_data():
    test_data = pd.read_csv('/data1/zmy/data/divide_csv/three/test.csv')
    train_data = pd.read_csv('/data1/zmy/data/divide_csv/three/train.csv')
    
    test_features = []
    test_labels = []
    test_ctslice = []
    test_petslice = []
    
    train_features = []
    train_labels = []
    train_ctslice = []
    train_petslice = []

    for i in range(len(test_data)):
        one_feature = [test_data['z'][i], test_data['x'][i]/512, test_data['y'][i]/512, test_data['r'][i],
                       test_data['patientWeight'][i], test_data['patientSex'][i], test_data['patientAge'][i],
                       test_data['patientSize'][i], test_data['local_suvmax'][i], test_data['local_suvmin'][i],
                       test_data['local_suvavg'][i], test_data['local_suvstd'][i], test_data['local_suvvar'][i]]

        test_features.append(one_feature)
        test_labels.append(int(test_data['cancer_type'][i]) - 1)
        ct_name = '/data1/zmy/data/Slice/three/test/CTSlice/'+str(test_data['patientID'][i])+'_'+str(test_data['cancer_type'][i])+'_CTSlice.npy'
        ct = np.load(ct_name)
        test_ctslice.append(ct)
        pet_name = '/data1/zmy/data/Slice/three/test/PETSlice/'+str(test_data['patientID'][i])+'_'+str(test_data['cancer_type'][i])+'_PETSlice.npy'
        pet = np.load(pet_name)
        test_petslice.append(pet)

    for j in range(len(train_data)):
        one_feature = [train_data['z'][j], train_data['x'][j]/512, train_data['y'][j]/512, train_data['r'][j],
                       train_data['patientWeight'][j], train_data['patientSex'][j], train_data['patientAge'][j],
                       train_data['patientSize'][j], train_data['local_suvmax'][j], train_data['local_suvmin'][j],
                       train_data['local_suvavg'][j],  train_data['local_suvstd'][j], train_data['local_suvvar'][j]]
        train_features.append(one_feature)
        train_labels.append(int(train_data['cancer_type'][j]) - 1)
        ct_name = '/data1/zmy/data/Slice/three/train/CTSlice/'+str(train_data['patientID'][j])+'_'+str(train_data['cancer_type'][j])+'_CTSlice.npy'
        ct = np.load(ct_name)
        train_ctslice.append(ct)
        pet_name = '/data1/zmy/data/Slice/three/train/PETSlice/'+str(train_data['patientID'][j])+'_'+str(train_data['cancer_type'][j])+'_PETSlice.npy'
        pet = np.load(pet_name)
        train_petslice.append(pet)

    x_train = np.asarray(train_features, dtype=np.float)
    x_train_ct = np.asarray(train_ctslice)
    x_train_pet = np.asarray(train_petslice)
    y_train = np.asarray(train_labels, dtype=np.int)

    x_test = np.asarray(test_features, dtype=np.float)
    x_test_ct = np.asarray(test_ctslice)
    x_test_pet = np.asarray(test_petslice)
    y_test = np.asarray(test_labels, dtype=np.int)

    x_train = standard_scaler(x_train)
    x_test = standard_scaler(x_test)

    # train test ct和pet进行合并
    x_train_img = np.stack((x_train_ct, x_train_pet), axis=3)
    x_test_img = np.stack((x_test_ct, x_test_pet), axis=3)

    # shuffle trainset
    m = x_train_img.shape[0]
    permutation = list(np.random.permutation(m))
    x_train_img = x_train_img[permutation, :, :, :]
    x_train = x_train[permutation, :]
    y_train = y_train[permutation]

    # to_category
    y_train = to_categorical(y_train, 3)
    y_test = to_categorical(y_test, 3)

    return x_train, x_train_img, y_train, x_test, x_test_img,y_test


def train():
    x_train, x_train_img, Y_train, x_test, x_test_img,Y_test = read_data()

    print("X_train shape: " + str(x_train.shape))
    print("X_train_img shape: " + str(x_train_img.shape))
    print("Y_train shape: " + str(Y_train.shape))
    print("X_test shape: " + str(x_test.shape))
    print("X_test_img shape: " + str(x_test_img.shape))
    print("Y_test shape: " + str(Y_test.shape))

    model = ResNet50()
    # adam = Adam(lr=Learn_rate)
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    model.summary()


    if not os.path.exists(work_dir):
        os.makedirs(work_dir)
    checkpoint = ModelCheckpoint(filepath=work_dir + "/" + model_name + "_" + "_e" + "{epoch:02d}-{val_loss:.4f}.hd5",
                                 monitor='val_loss',
                                 verbose=1,
                                 save_best_only=False,
                                 save_weights_only=False,
                                 mode='auto',
                                 period=1)
    # 每隔一轮且每当val_loss降低时保存一次模型
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    checkpoint_fixed_name = ModelCheckpoint(filepath=best_model_path,
                                            monitor='val_loss',
                                            verbose=1,
                                            save_best_only=True,
                                            save_weights_only=False,
                                            mode='auto',
                                            period=1)
    
    model.fit(x=[x_train_img, x_train],
              y=Y_train,
              batch_size=Batch_size,
              epochs=EPOCH,
              validation_data=([x_test_img, x_test], Y_test),
              callbacks=[checkpoint,
                         checkpoint_fixed_name,
                         TensorBoard(log_dir=log_dir)]
              )

def predict():
    """
    进行预测
    """

    x_train, x_train_img, Y_train, x_test, x_test_img,Y_test = read_data()

    # 模型
    model_path = "/home/zmy/workspace2021/three_models/two_resnet50_ct_pet_1/two_resnet50_ct_pet_1_best.hd5"
    model = ResNet50()
    model.load_weights(model_path, by_name=False)

    # 开始计时
    start_time = datetime.datetime.now()
    result = model.predict([x_train_img, x_train])

    # 测试花费时间
    current_time = datetime.datetime.now()
    res = current_time - start_time
    print("Done in : ", res.total_seconds(), " seconds")
    print("model: " + model_path)
    print("label:", Y_test)
    print("predict:", result)


    # 将三分类概率结果转换为下标标签
    preds = np.argmax(result, axis=1)
    y_labels = np.argmax(Y_test, axis=1)
    # print(preds)
    # print(y_test)

    # 计算评估指标
    target_names = ['1', '2', '3']
    result_statis = classification_report(y_labels, preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(y_labels, preds)
    print(confusion)


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    if not os.path.exists(model_save_path):
        os.makedirs(model_save_path)
    train()
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    shutil.copy(best_model_path, model_save_path)

